from itertools import chain
from collections import defaultdict
import math
import os
import re
from decimal import *
class naivebayes():
    def __init__(self):
        self.allwords_set=set()
        self.tests=[]
        self.spamwords=[]
        self.hamwords=[]
        self.target=[]
        self.allwords=[]
    def load_data(self):
        path1='spam'
        filenames1=os.listdir(path1)
        for filename in filenames1[:490]:
            f=open(path1+'/'+filename,'r',encoding='utf8',errors='ignore')
            for i in f:
                if i.startswith('Subject:'):
                    i=re.sub(r'[#,*,@,^,?,!,~,=,;,:,.,`,/,&,),(,+,|]','',i).strip().strip('/d').replace(']','').replace('[','')[8:].lower()
                    i=i.split()
                    self.spamwords.append(i)
            f.close()
        path2='ham'
        filenames2=os.listdir(path2)
        for filename in filenames2[:1000]:
            f=open(path2+'/'+filename,'r',encoding='utf8',errors='ignore')
            for i in f:
                if i.startswith('Subject:'):
                    i=re.sub(r'[#,*,@,^,?,!,~,=,;,:,.,+,`,/,&,),(,|]','',i).strip().replace(']','').replace('[','')[11:].lower().split()
                    self.hamwords.append(i)
            f.close()

    def fit(self,allwords,labels):
        self.target=labels
        self.train=allwords
        def factory():
            return 0
        self.spamwords_count=defaultdict(factory)
        self.hamwords_count=defaultdict(factory)
        for data,target in zip(self.train,self.target):
            if target==1:
                for spam in data:
                    self.spamwords_count[spam]+=1
            else:
                for ham in data:
                    self.hamwords_count[ham]+=1
        for spamword in list(chain(*self.spamwords)):
            self.allwords_set.add(spamword)
        for hamword in list(chain(*self.hamwords)):
            self.allwords_set.add(hamword) 
    def predict(self,tests):
        self.tests=tests
        ham_pro = 1
        spam_pro = 1
        for word in self.tests:
            spam_pro *= ((self.spamwords_count.get(word, 0) + 1) / len(list(chain(*self.spamwords))))
            ham_pro *= ((self.hamwords_count.get(word, 0) + 1) / len(list(chain(*self.hamwords))))
        spam_pro *= len(self.spamwords) / (len(self.hamwords) + len(self.spamwords))
        ham_pro *= len(self.hamwords) / (len(self.hamwords) + len(self.spamwords))
        spam_prob = spam_pro / (spam_pro + ham_pro)
        ham_prob = ham_pro / (spam_pro + ham_pro)
        return (spam_prob,ham_prob,int(spam_prob>ham_prob))

    def predict_prob(self):
        print('垃圾邮件的概率为{}，非垃圾邮件的概率为{}'.format(self.predict(self.tests)[0], self.predict(self.tests)[1]))
    def accuracy(self,tests,spam_lenth,ham_lenth):
        isspam=0
        isham=0
        for test in tests[:spam_lenth]:
            if self.predict(test)[2]==1:
                isspam+=1
        for test in tests[spam_lenth:spam_lenth+ham_lenth]:
            if self.predict(test)[2]==0:
                isham+=1
        getcontext=10
        print(Decimal(isspam+isham)/Decimal(spam_lenth+ham_lenth))
    
    
    
path1='spam'
filenames1=os.listdir(path1)
spamwords=[]
hamwords=[]
allwords=[]
for filename in filenames1[:450]:
    f=open(path1+'/'+filename,'r',encoding='utf8',errors='ignore')
    for i in f:
        if i.startswith('Subject:'):
            i=re.sub(r'[#,*,@,^,?,!,~,=,;,:,.,`,/,&,),(,+,|]','',i).strip().strip('/d').replace(']','').replace('[','')[8:].lower()
            i=i.split()
            spamwords.append(i)
    f.close()
path2='ham'
filenames2=os.listdir(path2)
for filename in filenames2[:500]:
    f=open(path2+'/'+filename,'r',encoding='utf8',errors='ignore')
    for i in f:
        if i.startswith('Subject:'):
            i=re.sub(r'[#,*,@,^,?,!,~,=,;,:,.,+,`,/,&,),(,|]','',i).strip().replace(']','').replace('[','')[11:].lower().split()
            hamwords.append(i)
    f.close()
labels=[1]*450+[0]*500
allwords=spamwords+hamwords

tests=[]
for filename in filenames1[:50]:
    f=open(path1+'/'+filename,'r',encoding='utf8',errors='ignore')
    for i in f:
        if i.startswith('Subject:'):
            i=re.sub(r'[#,*,@,^,?,!,~,=,;,:,.,+,`,/,&,),(,|]','',i).strip().replace(']','').replace('[','')[11:].lower().split()
            tests.append(i)
    f.close()
for filename in filenames2[300:400]:
    f=open(path2+'/'+filename,'r',encoding='utf8',errors='ignore')
    for i in f:
        if i.startswith('Subject:'):
            i=re.sub(r'[#,*,@,^,?,!,~,=,;,:,.,+,`,/,&,),(,|]','',i).strip().replace(']','').replace('[','')[11:].lower().split()
            tests.append(i)
    f.close()
    
a=naivebayes()
a.load_data()
a.fit(allwords,labels)
a.predict(['you', 'have', 'been', 'hand', 'selected', 'free', 'info'])
a.predict_prob()
a.predict(['the', 'gov', 'gets', 'tough', 'on', 'net', 'userser', 'pirates'])
a.predict_prob()
a.accuracy(tests,50,100)
